package com.forest.util;

import cn.hutool.core.codec.Base64;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.HexUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.KeyUtil;
import cn.hutool.crypto.SecureUtil;
import cn.hutool.json.JSONUtil;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import java.nio.ByteBuffer;
import java.security.Provider;
import java.security.SecureRandom;
import java.security.Security;

/**
 * 国密SM4加解密工具类(GCM模式)
 *
 * @author Administrator
 */
public class Sm4Util {

    private static final Logger LOG = LoggerFactory.getLogger(Sm4Util.class);
    /**
     * GCM模式下的认证标签长度，推荐为128位
     */
    private static final Integer GCM_AUTH_TAG_LENGTH_BYTES = 128;

    private static final String ALGORITHM = "SM4";
    private static final String PROVIDER = "BC";
    private static final String ALGORITHM_MODE_PADDING = "SM4/GCM/NoPadding";
    /** 混淆参数长度 */
    private static final Integer IV_LENGTH_BYTES = 12;

    /** 加密nonce长度 */
    private static final Integer NONCE_LENGTH_BYTES = 12;

    /** 密钥长度 */
    private static final Integer KEY_LENGTH = 16;

    static {
        try {
            //初始化BouncyCastleProvider
            Provider provider = Security.getProvider(PROVIDER);
            if (provider == null) {
                Security.addProvider(new BouncyCastleProvider());
            }
        } catch (Exception e) {
            LOG.error("初始化失败: {}", e.getMessage(), e);
            if(EnvironmentHolder.getUseBouncyCastle() && EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("初始化国密提供者失败");
            }
        }
    }

    public static byte[] encrypt(String data, String keyStr) {
        return encrypt(StrUtil.bytes(data, CharsetUtil.CHARSET_UTF_8), keyStr, null, true);
    }

    public static byte[] encrypt(byte[] data, String keyStr, String nonceStr, boolean wrap) {
        String secretKeyStr = StrUtil.isBlank(keyStr) ? EnvironmentHolder.getSecretKey() : keyStr;
        if(StrUtil.isBlank(secretKeyStr)) {
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("请设置密钥");
            } else {
                return data;
            }
        }
        byte[] key = SecureUtil.decode(secretKeyStr);
        if(key.length != KEY_LENGTH) {
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("请设置密钥");
            } else {
                return data;
            }
        }
        //创建key对象
        SecretKey secretKey = KeyUtil.generateKey(ALGORITHM, key);
        byte [] nonce =  StrUtil.isBlank(nonceStr) ? randomBytes(NONCE_LENGTH_BYTES) : SecureUtil.decode(nonceStr);
        //创建加密参数
        GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(GCM_AUTH_TAG_LENGTH_BYTES, nonce);
        try {
            Cipher cipher = Cipher.getInstance(ALGORITHM_MODE_PADDING, PROVIDER);
            cipher.init(Cipher.ENCRYPT_MODE, secretKey, gcmParameterSpec);
            byte[] encrypted = cipher.doFinal(data);
            return wrap ? wrap(nonce, encrypted, StrUtil.isNotBlank(nonceStr)) : encrypted ;
        } catch (Exception e) {
            LOG.error("加密失败: {}", e.getMessage(), e);
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("加密失败: " + e.getMessage());
            } else {
                return data;
            }
        }
    }

    /**
     * 封装数据
     * @param nonce
     * @param encrypted
     * @return
     */
    private static byte[] wrap(byte[] nonce, byte[] encrypted, boolean fixedIv) {
        ByteBuffer buffer;
        if(fixedIv) {
            buffer = ByteBuffer.allocate(NONCE_LENGTH_BYTES + encrypted.length);
            buffer.put(nonce, 0, NONCE_LENGTH_BYTES/2);
            buffer.put(encrypted);
            buffer.put(nonce, NONCE_LENGTH_BYTES/2, NONCE_LENGTH_BYTES/2);
        } else {
            byte[] iv = randomBytes(IV_LENGTH_BYTES);
            buffer = ByteBuffer.allocate(NONCE_LENGTH_BYTES + IV_LENGTH_BYTES + encrypted.length);
            buffer.put(nonce, 0, NONCE_LENGTH_BYTES/2);
            buffer.put(iv,0, IV_LENGTH_BYTES/2);
            buffer.put(encrypted);
            buffer.put(nonce, NONCE_LENGTH_BYTES/2, NONCE_LENGTH_BYTES/2);
            buffer.put(iv, IV_LENGTH_BYTES/2, IV_LENGTH_BYTES/2);
        }
        return buffer.array();
    }

    public static byte[] randomBytes(int length) {
        byte[] bytes = new byte[length];
        SecureRandom random = new SecureRandom();
        random.nextBytes(bytes);
        return bytes;
    }

    public static String encryptBase64(String data) {
        byte[] encrypt = encrypt(data, null);
        return Base64.encode(encrypt);
    }

    public static String encryptBase64(String data, String key) {
        byte[] encrypt = encrypt(data, key);
        return Base64.encode(encrypt);
    }

    public static String encryptBase64(String data, String key, String iv) {
        byte[] encrypt = encrypt(StrUtil.bytes(data, CharsetUtil.CHARSET_UTF_8), key, iv, true);
        return Base64.encode(encrypt);
    }

    public static String encryptBase64(String data, String key, String iv, boolean wrap) {
        byte[] encrypt = encrypt(StrUtil.bytes(data, CharsetUtil.CHARSET_UTF_8), key, iv, wrap);
        return Base64.encode(encrypt);
    }

    public static String encryptHex(String data) {
        byte [] encrypt = encrypt(data, null);
        return HexUtil.encodeHexStr(encrypt);
    }

    public static String encryptHex(String data, String key) {
        byte [] encrypt = encrypt(data, key);
        return HexUtil.encodeHexStr(encrypt);
    }

    public static String encryptHex(String data, String key, String iv) {
        byte[] encrypt = encrypt(StrUtil.bytes(data, CharsetUtil.CHARSET_UTF_8), key, iv, true);
        return HexUtil.encodeHexStr(encrypt);
    }

    public static String encryptHex(String data, String key, String iv, boolean wrap) {
        byte[] encrypt = encrypt(StrUtil.bytes(data, CharsetUtil.CHARSET_UTF_8), key, iv, wrap);
        return HexUtil.encodeHexStr(encrypt);
    }

    /**
     * 解密.
     *
     * @param decode
     * @return
     */
    public static byte[] decrypt(byte[] decode, String keyStr, String nonceStr, boolean wrap) {
        String secretKeyStr = StrUtil.isBlank(keyStr) ? EnvironmentHolder.getSecretKey() : keyStr;
        if(StrUtil.isBlank(secretKeyStr)) {
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("请设置密钥");
            } else {
                return decode;
            }
        }
        byte[] key = SecureUtil.decode(secretKeyStr);
        if(key.length != KEY_LENGTH) {
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("请设置密钥");
            } else {
                return decode;
            }
        }
        //创建key对象
        SecretKey secretKey = KeyUtil.generateKey(ALGORITHM, key);
        //创建加密参数
        byte[] nonce = StrUtil.isBlank(nonceStr) ? unwrapNonce(decode, StrUtil.isNotBlank(nonceStr)) : SecureUtil.decode(nonceStr);
        //解密数据
        byte[] encrypted = unwrapEncrypted(decode, StrUtil.isNotBlank(nonceStr), wrap);
        //创建加密参数
        GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(GCM_AUTH_TAG_LENGTH_BYTES, nonce);
        try {
            Cipher cipher = Cipher.getInstance(ALGORITHM_MODE_PADDING, PROVIDER);
            cipher.init(Cipher.DECRYPT_MODE, secretKey, gcmParameterSpec);
            byte[] decrypt = cipher.doFinal(encrypted);
            return decrypt;
        } catch (Exception e) {
            LOG.error("解密失败: {}", e.getMessage(), e);
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("加密失败: " + e.getMessage());
            } else {
                return decode;
            }
        }
    }

    private static byte[] unwrapEncrypted(byte[] data, boolean fixedIv, boolean wrap) {
        //未封装数据，整个data都是密文
        if(!wrap) {
            return data;
        }
        int middleNonceIndex = NONCE_LENGTH_BYTES / 2;
        int middleIvIndex = IV_LENGTH_BYTES / 2;
        int dataLength = fixedIv ? data.length - NONCE_LENGTH_BYTES : data.length - NONCE_LENGTH_BYTES - IV_LENGTH_BYTES;
        int dataStartIndex = !fixedIv ? middleNonceIndex + middleIvIndex : middleNonceIndex;
        byte[] encrypted = new byte[dataLength];
        System.arraycopy(data, dataStartIndex , encrypted, 0, encrypted.length);
        return encrypted;
    }

    private static byte[] unwrapNonce(byte[] data, boolean fixedIv) {
        int middleNonceIndex = NONCE_LENGTH_BYTES / 2;
        byte [] nonce = new byte[NONCE_LENGTH_BYTES];
        System.arraycopy(data, 0, nonce, 0, middleNonceIndex);
        System.arraycopy(data,  data.length - (fixedIv ? middleNonceIndex : ((NONCE_LENGTH_BYTES + IV_LENGTH_BYTES)/2)), nonce, middleNonceIndex,  middleNonceIndex);
        return nonce;
    }

    /**
     * 解密(hex/base64)
     * @param data
     * @return
     */
    public static String decrypt(String data) {
        byte[] decode = SecureUtil.decode(data);
        byte[] decrypt = decrypt(decode, null,  null, true);
        return decrypt == null ? data : new String(decrypt);
    }

    public static String decrypt(String data, String key) {
        byte[] decode = SecureUtil.decode(data);
        byte[] decrypt = decrypt(decode, key,  null, true);
        return decrypt == null ? data : new String(decrypt);
    }

    public static String decrypt(String data, String key, String iv) {
        byte[] decode = SecureUtil.decode(data);
        byte[] decrypt = decrypt(decode, key, iv, true);
        return decrypt == null ? data : new String(decrypt);
    }

    public static String decrypt(String data, String key, String iv, boolean wrap) {
        byte[] decode = SecureUtil.decode(data);
        byte[] decrypt = decrypt(decode, key, iv, wrap);
        return decrypt == null ? data : new String(decrypt);
    }
}
